#ifndef AES_H
#define AES_H

#include <stdint.h>


// #define the macros below to 1/0 to enable/disable the mode of operation.
//
// CBC enables AES encryption in CBC-mode of operation.
// CTR enables encryption in counter-mode.
// ECB enables the basic ECB 16-byte block algorithm. All can be enabled simultaneously.

// The #ifndef-guard allows it to be configured before #include'ing or at compile time.
#ifndef CBC
  #define CBC 1
#endif

#ifndef ECB
  #define ECB 0
#endif

#ifndef CTR
  #define CTR 0
#endif


//#define AES128 1
//#define AES192 1
#define AES256                  1

#define AES_BLOCKLEN            16      // Block length in bytes - AES is 128b block only

#if defined(AES256) && (AES256 == 1)
    #define AES_KEYLEN          32
    #define AES_keyExpSize      240
#elif defined(AES192) && (AES192 == 1)
    #define AES_KEYLEN          24
    #define AES_keyExpSize      208
#else
    #define AES_KEYLEN          16      // Key length in bytes
    #define AES_keyExpSize      176
#endif


//struct AES_ctx
//{
//    uint8_t RoundKey[AES_keyExpSize];
//#if (defined(CBC) && (CBC == 1)) || (defined(CTR) && (CTR == 1))
//    uint8_t Iv[AES_BLOCKLEN];
//#endif
//};


class Aes
{
public:
    Aes();
    Aes(const uint8_t *key, const uint8_t *iv);


    void init(const uint8_t* key);
#if (defined(CBC) && (CBC == 1)) || (defined(CTR) && (CTR == 1))
    void initKeyIV(const uint8_t* key, const uint8_t* iv);
    void setIV(const uint8_t* iv);
#endif

#if defined(ECB) && (ECB == 1)
    // buffer size is exactly AES_BLOCKLEN bytes;
    // you need only AES_init_ctx as IV is not used in ECB
    // NB: ECB is considered insecure for most uses
    void ECB_encrypt(uint8_t* buf);
    void ECB_decrypt(uint8_t* buf);
#endif // #if defined(ECB) && (ECB == !)

#if defined(CBC) && (CBC == 1)
    // buffer size MUST be mutile of AES_BLOCKLEN;
    // Suggest https://en.wikipedia.org/wiki/Padding_(cryptography)#PKCS7 for padding scheme
    // NOTES: you need to set IV in ctx via AES_init_ctx_iv() or AES_ctx_set_iv()
    //        no IV should ever be reused with the same key
    void CBC_encrypt_buffer(uint8_t* buf, uint32_t length);
    void CBC_decrypt_buffer(uint8_t* buf, uint32_t length);

#endif // #if defined(CBC) && (CBC == 1)

#if defined(CTR) && (CTR == 1)

    // Same function for encrypting as for decrypting.
    // IV is incremented for every block, and used after encryption as XOR-compliment for output
    // Suggesting https://en.wikipedia.org/wiki/Padding_(cryptography)#PKCS7 for padding scheme
    // NOTES: you need to set IV in ctx with AES_init_ctx_iv() or AES_ctx_set_iv()
    //        no IV should ever be reused with the same key
    void CTR_xcrypt_buffer(uint8_t* buf, uint32_t length);

#endif // #if defined(CTR) && (CTR == 1)

    void showKeyIV(void);

private:
    uint8_t RoundKey[AES_keyExpSize];
#if (defined(CBC) && (CBC == 1)) || (defined(CTR) && (CTR == 1))
    uint8_t Iv[AES_BLOCKLEN];
#endif
};

#endif // AES_H
